package com.zengde.ai_sandbox.controller;

import com.zengde.ai_sandbox.manager.BigModelManager;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
import org.springframework.context.annotation.Profile;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.Payload;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Controller;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Controller
@ConditionalOnWebApplication(type = ConditionalOnWebApplication.Type.SERVLET)
@Profile("!test")
@RequiredArgsConstructor
@Slf4j
public class ChatWebSocketController {
    
    private final BigModelManager bigModelManager;
    private final SimpMessagingTemplate messagingTemplate;
    
    // 管理活动的会话订阅，用于会话关闭时取消订阅
    private final Map<String, Flux<?>> activeSessionFluxes = new ConcurrentHashMap<>();
    
    /**
     * 处理客服发送的文本消息
     */
    @MessageMapping("/chat.sendMessage")
    public Mono<Void> handleCustomerServiceMessage(@Payload Map<String, String> message) {
        log.info("收到发送消息请求: {}", message);
        String sessionId = message.get("sessionId");
        String content = message.get("content");
        String customerId = message.get("customerId");
        String modelType = message.getOrDefault("modelType", "normal"); // normal, think
        String responseType = message.getOrDefault("responseType", "stream"); // stream, sync
        
        log.info("处理客服消息: sessionId={}, content={}, modelType={}, responseType={}", 
                sessionId, content, modelType, responseType);
        
        try {
            return processTextMessage(sessionId, customerId, content, modelType, responseType);
        } catch (Exception e) {
            log.error("处理消息时发生异常", e);
            sendError(sessionId, customerId, e);
            return Mono.empty();
        }
    }
    
    /**
     * 处理客服发送的图片消息
     */
    @MessageMapping("/chat.sendImageMessage")
    public Mono<Void> handleImageMessage(@Payload Map<String, String> message) {
        log.info("收到图片消息请求: {}", message);
        String sessionId = message.get("sessionId");
        String content = message.get("content");
        String imageUrl = message.get("imageUrl");
        String customerId = message.get("customerId");
        
        log.info("处理图片消息: sessionId={}, content={}, imageUrl={}", 
                sessionId, content, imageUrl);
        
        try {
            if (imageUrl == null || imageUrl.isEmpty()) {
                throw new IllegalArgumentException("图片URL不能为空");
            }
            
            return Mono.fromCallable(() -> {
                String response = bigModelManager.imgChat(sessionId, content, imageUrl);
                log.info("获取到图片消息响应: {}", response);
                sendFullResponse(sessionId, customerId, response);
                return response;
            }).then();
        } catch (Exception e) {
            log.error("处理图片消息时发生异常", e);
            sendError(sessionId, customerId, e);
            return Mono.empty();
        }
    }
    
    /**
     * 处理文本消息的不同调用方式
     */
    private Mono<Void> processTextMessage(String sessionId, String customerId, String content, 
                                    String modelType, String responseType) {
        boolean isThinkingMode = "think".equals(modelType);
        boolean isStreamMode = "stream".equals(responseType);
        
        if (isThinkingMode) {
            if (isStreamMode) {
                log.info("使用深度思考流式调用");
                // 创建Flux流并应用背压和限流策略
                Flux<String> responseFlux = bigModelManager.ts(sessionId, content)
                        .doOnNext(chunk -> sendResponseChunk(sessionId, customerId, chunk))
                        .doOnError(error -> sendError(sessionId, customerId, error))
                        .doOnComplete(() -> sendComplete(sessionId, customerId))
                        .doFinally(signalType -> {
                            if (signalType == SignalType.ON_COMPLETE || signalType == SignalType.ON_ERROR) {
                                activeSessionFluxes.remove(sessionId);
                                log.info("流式调用完成或出错，移除会话流: sessionId={}, signalType={}", sessionId, signalType);
                            }
                        })
                        .onBackpressureBuffer(256, dropped -> log.warn("背压缓冲区溢出，丢弃消息: {}", dropped))
                        .limitRate(32);
                
                // 保存活动的Flux，以便后续可以取消
                activeSessionFluxes.put(sessionId, responseFlux);
                
                // 返回Mono，让Spring管理订阅生命周期
                return responseFlux.then();
            } else {
                log.info("使用深度思考非流式调用");
                return Mono.fromCallable(() -> {
                    String response = bigModelManager.tns(sessionId, content);
                    log.info("获取到深度思考非流式响应: {}", response);
                    sendFullResponse(sessionId, customerId, response);
                    return response;
                }).then();
            }
        } else {
            if (isStreamMode) {
                log.info("使用非深度思考流式调用");
                // 创建Flux流并应用背压和限流策略
                Flux<String> responseFlux = bigModelManager.nts(sessionId, content)
                        .doOnNext(chunk -> sendResponseChunk(sessionId, customerId, chunk))
                        .doOnError(error -> sendError(sessionId, customerId, error))
                        .doOnComplete(() -> sendComplete(sessionId, customerId))
                        .doFinally(signalType -> {
                            if (signalType == SignalType.ON_COMPLETE || signalType == SignalType.ON_ERROR) {
                                activeSessionFluxes.remove(sessionId);
                                log.info("流式调用完成或出错，移除会话流: sessionId={}, signalType={}", sessionId, signalType);
                            }
                        })
                        .onBackpressureBuffer(256, dropped -> log.warn("背压缓冲区溢出，丢弃消息: {}", dropped))
                        .limitRate(32);
                
                // 保存活动的Flux，以便后续可以取消
                activeSessionFluxes.put(sessionId, responseFlux);
                
                // 返回Mono，让Spring管理订阅生命周期
                return responseFlux.then();
            } else {
                log.info("使用非深度思考非流式调用");
                return Mono.fromCallable(() -> {
                    String response = bigModelManager.ntns(sessionId, content);
                    log.info("获取到非深度思考非流式响应: {}", response);
                    sendFullResponse(sessionId, customerId, response);
                    return response;
                }).then();
            }
        }
    }
    
    /**
     * 创建新会话
     */
    @MessageMapping("/chat.createSession")
    public Mono<Void> createSession(@Payload Map<String, String> request) {
        log.info("收到创建会话请求: {}", request);
        String customerId = request.get("customerId");
        String modelServiceName = request.getOrDefault("modelServiceName", "doubaoService");
        
        try {
            if (customerId == null || customerId.isEmpty()) {
                throw new IllegalArgumentException("客服ID不能为空");
            }
            
            log.info("开始创建会话: customerId={}, modelServiceName={}", customerId, modelServiceName);
            
            return Mono.fromCallable(() -> {
                String sessionId = bigModelManager.createSession(customerId, modelServiceName);
                log.info("会话创建成功: sessionId={}", sessionId);
                
                // 获取机器人的首次发言
                log.info("获取机器人首次发言");
                String firstMessage = bigModelManager.robotFirstMessage(sessionId);
                log.info("获取到机器人首次发言: {}", firstMessage);
                
                Map<String, Object> response = new HashMap<>();
                response.put("type", "SESSION_CREATED");
                response.put("sessionId", sessionId);
                response.put("firstMessage", firstMessage);
                
                log.info("发送会话创建成功消息: {}", response);
                messagingTemplate.convertAndSendToUser(
                        customerId, 
                        "/queue/chat", 
                        response
                );
                log.info("会话创建成功消息已发送");
                
                return sessionId;
            }).onErrorResume(e -> {
                log.error("创建会话时发生异常", e);
                Map<String, Object> errorResponse = new HashMap<>();
                errorResponse.put("type", "ERROR");
                errorResponse.put("error", e.getMessage());
                
                messagingTemplate.convertAndSendToUser(
                        customerId, 
                        "/queue/chat", 
                        errorResponse
                );
                return Mono.empty();
            }).then();
        } catch (Exception e) {
            log.error("创建会话时发生异常", e);
            Map<String, Object> errorResponse = new HashMap<>();
            errorResponse.put("type", "ERROR");
            errorResponse.put("error", e.getMessage());
            
            messagingTemplate.convertAndSendToUser(
                    customerId, 
                    "/queue/chat", 
                    errorResponse
            );
            return Mono.empty();
        }
    }
    
    /**
     * 关闭会话
     */
    @MessageMapping("/chat.closeSession")
    public Mono<Void> closeSession(@Payload Map<String, String> request) {
        log.info("收到关闭会话请求: {}", request);
        String sessionId = request.get("sessionId");
        String customerId = request.get("customerId");
        
        try {
            if (sessionId == null || sessionId.isEmpty()) {
                throw new IllegalArgumentException("会话ID不能为空");
            }
            
            return Mono.fromRunnable(() -> {
                log.info("关闭会话: sessionId={}", sessionId);
                
                // 取消正在进行的流式调用（如果有）
                Flux<?> sessionFlux = activeSessionFluxes.remove(sessionId);
                if (sessionFlux != null) {
                    log.info("取消会话的活动流: sessionId={}", sessionId);
                    // 不需要额外操作，移除Map中的引用后，Flux会被垃圾回收
                }
                
                // 从BigModelManager中移除会话
                bigModelManager.removeSession(sessionId);
                
                Map<String, Object> response = new HashMap<>();
                response.put("type", "SESSION_CLOSED");
                response.put("sessionId", sessionId);
                
                log.info("发送会话关闭消息: {}", response);
                messagingTemplate.convertAndSendToUser(
                        customerId, 
                        "/queue/chat", 
                        response
                );
                log.info("会话关闭消息已发送");
            }).onErrorResume(e -> {
                log.error("关闭会话时发生异常", e);
                sendError(sessionId, customerId, e);
                return Mono.empty();
            }).then();
        } catch (Exception e) {
            log.error("关闭会话时发生异常", e);
            sendError(sessionId, customerId, e);
            return Mono.empty();
        }
    }
    
    /**
     * 发送响应片段
     */
    private void sendResponseChunk(String sessionId, String customerId, String chunk) {
        Map<String, Object> response = new HashMap<>();
        response.put("type", "RESPONSE_CHUNK");
        response.put("sessionId", sessionId);
        response.put("content", chunk);
        
        log.debug("发送响应片段: sessionId={}, customerId={}, chunk={}", sessionId, customerId, chunk);
        messagingTemplate.convertAndSendToUser(
                customerId, 
                "/queue/chat", 
                response
        );
    }
    
    /**
     * 发送完整响应
     */
    private void sendFullResponse(String sessionId, String customerId, String content) {
        Map<String, Object> response = new HashMap<>();
        response.put("type", "FULL_RESPONSE");
        response.put("sessionId", sessionId);
        response.put("content", content);
        
        log.info("发送完整响应: sessionId={}, customerId={}, content长度={}", 
                sessionId, customerId, content != null ? content.length() : 0);
        messagingTemplate.convertAndSendToUser(
                customerId, 
                "/queue/chat", 
                response
        );
        log.info("完整响应已发送");
    }
    
    /**
     * 发送完成消息
     */
    private void sendComplete(String sessionId, String customerId) {
        Map<String, Object> response = new HashMap<>();
        response.put("type", "RESPONSE_COMPLETE");
        response.put("sessionId", sessionId);
        
        log.info("发送完成消息: sessionId={}, customerId={}", sessionId, customerId);
        messagingTemplate.convertAndSendToUser(
                customerId, 
                "/queue/chat", 
                response
        );
        log.info("完成消息已发送");
    }
    
    /**
     * 发送错误消息
     */
    private void sendError(String sessionId, String customerId, Throwable error) {
        Map<String, Object> response = new HashMap<>();
        response.put("type", "ERROR");
        response.put("sessionId", sessionId);
        response.put("error", error.getMessage());
        
        log.error("发生错误: sessionId={}, customerId={}, error={}", sessionId, customerId, error.getMessage(), error);
        messagingTemplate.convertAndSendToUser(
                customerId, 
                "/queue/chat", 
                response
        );
        log.info("错误消息已发送");
    }
    
    /**
     * 获取会话历史消息
     */
    @MessageMapping("/chat.getHistory")
    public Mono<Void> getSessionHistory(@Payload Map<String, String> request) {
        log.info("收到获取历史消息请求: {}", request);
        String sessionId = request.get("sessionId");
        String customerId = request.get("customerId");
        
        try {
            return Mono.fromRunnable(() -> {
                // TODO: 实现从数据库获取历史消息的功能
                // 目前返回空的历史消息列表
                Map<String, Object> response = new HashMap<>();
                response.put("type", "HISTORY_MESSAGES");
                response.put("sessionId", sessionId);
                response.put("messages", new Object[]{}); // 空历史消息
                
                log.info("发送历史消息: {}", response);
                messagingTemplate.convertAndSendToUser(
                        customerId, 
                        "/queue/chat", 
                        response
                );
                log.info("历史消息已发送");
            }).onErrorResume(e -> {
                log.error("获取历史消息时发生异常", e);
                sendError(sessionId, customerId, e);
                return Mono.empty();
            }).then();
        } catch (Exception e) {
            log.error("获取历史消息时发生异常", e);
            sendError(sessionId, customerId, e);
            return Mono.empty();
        }
    }
    
    /**
     * 测试WebSocket通信
     */
    @MessageMapping("/chat.test")
    public Mono<Void> testWebSocket(@Payload Map<String, String> request) {
        log.info("收到测试请求: {}", request);
        String customerId = request.get("customerId");
        
        try {
            return Mono.fromRunnable(() -> {
                Map<String, Object> response = new HashMap<>();
                response.put("type", "TEST_RESPONSE");
                response.put("message", "这是一个测试响应");
                response.put("timestamp", System.currentTimeMillis());
                
                log.info("发送测试响应: {}", response);
                messagingTemplate.convertAndSendToUser(
                        customerId, 
                        "/queue/chat", 
                        response
                );
                log.info("测试响应已发送");
            }).then();
        } catch (Exception e) {
            log.error("测试WebSocket时发生异常", e);
            return Mono.empty();
        }
    }
} 